#!/usr/bin/env python3
"""
Point Cloud Saving Utility

This script saves point cloud data to .pcd format from unified data format.
It uses Open3D library for efficient and reliable PCD file writing.

Usage:
    This module provides savePcd function for saving point cloud data to .pcd files
    within the unified data format pipeline.
"""
import numpy as np
import os
import logging
import open3d as o3d
from scalePcd import scalePcd
from time_it import timeit
# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def ensure_dir(file_path):
    """Ensure the directory exists for the given file path."""
    directory = os.path.dirname(file_path)
    if directory and not os.path.exists(directory):
        os.makedirs(directory, exist_ok=True)


@timeit()
def savePcd(unifiedData: dict[str, np.ndarray], output_dir: str = "pcd_output", 
           filename_without_ext:str="pointcloud", typestr: str = "pointcloud", color: bool = True) -> dict[str, np.ndarray]:
    """
    Save point cloud data to .pcd file format from unified data format using Open3D.
    
    Args:
        unifiedData: dict containing single frame data with keys:
                    ["rgb":np.ndarray,"depth":np.ndarray,"pointcloud":np.ndarray,
                     "observer":np.ndarray,"endpose":np.ndarray,"qpos":np.ndarray]
                    where pointcloud is shape (Np, 6) for single frame [x, y, z, r, g, b]
        output_dir: Directory to save .pcd files
        episode_idx: Episode index for filename
        frame_idx: Frame index for filename
        color: Whether to save color information
        
    Returns:
        Dictionary containing original data (unchanged)
    """
    point_cloud_data = unifiedData[typestr]  # Shape: (Np, 6)

    
    if point_cloud_data is None:
        logger.warning("No point cloud data to save")
        return unifiedData
    
    # Validate input shape
    if len(point_cloud_data.shape) != 2 or point_cloud_data.shape[1] != 6:
        logger.warning(f"Expected point cloud shape (Np, 6), got {point_cloud_data.shape}. Skipping save.")
        return unifiedData
    
    # Filter out zero/invalid points
    valid_mask = np.any(point_cloud_data[:, :3] != 0, axis=1)
    
    if not np.any(valid_mask):
        logger.warning("No valid points in frame. Skipping save.")
        return unifiedData
    
    valid_points = point_cloud_data[valid_mask]
    
    # Create filename and full path
    filename = f"{filename_without_ext}.pcd"
    save_path = os.path.join(output_dir, filename)
    
    try:
        # Use the same logic as the existing save_pcd function
        # print(f"save_pcd {os.getcwd()} {save_path} {valid_points.shape}")
        ensure_dir(save_path)
        
        # Create Open3D point cloud object
        point_cloud = o3d.geometry.PointCloud()
        
        # Extract XYZ coordinates
        point_arr = valid_points[:, :3]
        point_cloud.points = o3d.utility.Vector3dVector(point_arr)
        
        # Add colors if requested and available
        if color and valid_points.shape[1] >= 6:
            colors_arr = valid_points[:, 3:6]  # RGB values
            
            # 使用 scalePcd 统一处理颜色缩放
            colors_scaled = scalePcd(colors_arr, tag="savePcd_colors")
            
            # Open3D expects colors in 0-1 range, so normalize for saving
            colors_normalized = colors_scaled / 255.0
            point_cloud.colors = o3d.utility.Vector3dVector(colors_normalized)
        
        # Save the point cloud
        success = o3d.io.write_point_cloud(save_path, point_cloud)
        
        if not success:
            logger.error(f"Failed to save point cloud to {save_path}")
        
    except Exception as e:
        logger.error(f"Error saving PCD file {save_path}: {e}")
    
    # Return original data unchanged
    return unifiedData 